#include "memory.h"
#include "stdint.h"
#include "print.h"
#include "debug.h"

#define PG_SIZE 4096

// ================ 位图地址 ================
// 为什么是 0xc009a000?
// 因为 0xc009f000 - 0x1000 - 0x1 = 0xc009e000 这是 PCB 表尾
// 0xc009e000 是将来主线程存放 PCB 的地址（表头 + 4k 为 0xc009efff）
// 因为页单位为 4k，位图以字节为单位划分，细粒度单位是位，每位表示 4k 即一页
// 故而一页大小的位图可管理的内存容量：4k * 8 * 4k = 128M
// 当前我们给虚拟机分配了 32Mb 内存空间，因此这才用了四分之一页
// 为了扩展，本书中这里假设要管理 4页 的位图，即最大可管理 512MB 的物理内存
// 综上所述，0xc009e000 还需要减去 4页 的位图：0xc009e000 - 0x4000 = 0xc009a000
#define MEM_BITMAP_BASE 0xc009a000
// =========================================

#define PDT_IDX(addr) ((addr & 0xffc00000) >> 22) // 获取 addr 的高 10 位
#define PTE_IDX(addr) ((addr & 0x003ff000) >> 12) // 获取 addr 的中间 10 位

/* 0xc0000000是内核从虚拟地址3G起. 0x100000意指跨过低端1M内存,使虚拟地址在逻辑上连续 */
#define K_HEAP_START 0xc0100000

// 内存池结构
struct pool {
    struct bitmap pool_bitmap; // 该内存池所用到的位图结构，用于管理物理内存
    uint32_t phy_addr_start;   // 该内存池所管理物理内存的起始地址
    uint32_t pool_size;        // 该内存池字节容量
};

struct pool kernel_pool, user_pool; // 内核物理内存池和用户物理内存池
struct virtual_addr kernel_vaddr;   // 用于给内核分配虚拟地址

// 在 pf 所表示的内存池中申请 pg_cnt 个虚拟页
// 成功时返回虚拟页的起始地址，失败则返回 NULL
static void* vaddr_get(enum pool_flags pf, uint32_t pg_cnt) {
    int vaddr_start = 0, bit_idx_start = -1;
    uint32_t cnt = 0;
    if(pf == PF_KERNEL) { // 内核内存池
        bit_idx_start = bitmap_scan(&kernel_vaddr.vaddr_bitmap, pg_cnt); // 扫描是否有足够的空间
        if(bit_idx_start == -1) { // 没有足够的空间
            return NULL;
        }
        while(cnt < pg_cnt) { // 有足够的空间，则将这 pg_cnt 个页的状态设置为 1 表示已占用
            bitmap_set(&kernel_vaddr.vaddr_bitmap, bit_idx_start + cnt++, 1);
        }

        // 返回的虚拟地址：内核虚拟内存池的起始地址 + 相对于内核内存池的偏移量 bit_idx_start * PG_SIZE(每位代表一页)
        vaddr_start = kernel_vaddr.vaddr_start + bit_idx_start * PG_SIZE;
    } else { // 用户内存池
        // ...
    }
    return (void*) vaddr_start;
}

// 得到虚拟地址 vaddr 对应的 PTE 指针
uint32_t* pte_ptr(uint32_t vaddr) {
    // 访问页目录的最后一个 PDE，将当前页目录表当成页表
    // 利用 vaddr 的高 10 位，去找到真正的页表的物理地址（按处理器的思想，这步得到的应该是物理页的物理地址，但这里我们将其看成真正的页表物理地址）
    // 利用 vaddr 的中间 10 位，作为页表的偏移量，从而找到 vaddr 对应的 PTE
    uint32_t* pte = (uint32_t*)(0xffc00000 + \
    ((vaddr & 0xffc00000) >> 10) + \
    PTE_IDX(vaddr) * 4); 
    return pte;
}

// 得到虚拟地址 vaddr 对应的 PDT 指针
uint32_t* pde_ptr(uint32_t vaddr) {
    uint32_t* pde = (uint32_t*) ((0xfffff000) + PDT_IDX(vaddr) * 4);
    return pde;
}

// 在 m_pool 指向的物理内存池中分配一个物理页
// 返回页框的物理地址，否则返回 NULL
static void* palloc(struct pool* m_pool) {
    // 扫描或设置位图要保证原子操作
    int bit_idx = bitmap_scan(&m_pool -> pool_bitmap, 1); // 找到一个物理页
    if(bit_idx == -1) return NULL;
    bitmap_set(&m_pool -> pool_bitmap, bit_idx, 1); // 设置该位使用状态为：已使用
    uint32_t page_phyaddr = ((bit_idx * PG_SIZE) + m_pool -> phy_addr_start);
    return (void*) page_phyaddr;
}

// 页表中添加虚拟地址 _vaddr 与物理地址 _page_phyaddr 的映射
static void page_table_add(void* _vaddr, void* _page_phyaddr) {
    uint32_t vaddr = (uint32_t) _vaddr, page_phyaddr = (uint32_t) _page_phyaddr;
    uint32_t* pde = pde_ptr(vaddr);
    uint32_t* pte = pte_ptr(vaddr);

    // 判断页目录项中的 P 位是否为 1，即判断 PDE 是否存在
    if(*pde & 0x00000001) { // 存在
        ASSERT(!(*pte & 0x00000001));
        // 判断 PTE 是否存在
        if(!(*pte & 0x00000001)) { // 不存在
            *pte = (page_phyaddr | PG_US_U | PG_RW_W | PG_P_1);
        } else {
            // 这个分支，我个人感觉压根没存在的必要
            // 因为 PTE 是不存在才创建，而这个分支表示已存在，那还创建个屁？
            PANIC("pte repeat");
            *pte = (page_phyaddr | PG_US_U | PG_RW_W | PG_P_1);
        }
    } else { // 不存在
        // 页表中用到的页框一律从内核内存池分配
        // PDE 不存在，自然 PTE 页不存在则需要创建一个新的页表
        uint32_t pde_phyaddr = (uint32_t) palloc(&kernel_pool);

        *pde = (pde_phyaddr | PG_US_U | PG_RW_W | PG_P_1);

        memset((void*)((int)pte & 0xfffff000), 0, PG_SIZE); // 避免脏读

        ASSERT(!(*pte & 0x00000001));

        *pte = (page_phyaddr | PG_US_U | PG_RW_W | PG_P_1);
    }
}

// 分配 pg_cnt 个页空间，返回虚拟地址，否则 NULL
void* malloc_page(enum pool_flags pf, uint32_t pg_cnt) {
    ASSERT(pg_cnt > 0 && pg_cnt < 3840);
    void* vaddr_start = vaddr_get(pf, pg_cnt);
    if(vaddr_start == NULL) return NULL;
    uint32_t vaddr = (uint32_t) vaddr_start, cnt = pg_cnt;
    struct pool* mem_pool = pf & PF_KERNEL ? &kernel_pool : &user_pool;

    // 因为虚拟地址是连续的，但物理地址可以是不连续的，所以要逐个做映射
    while(cnt-- > 0) {
        void* page_phyaddr = palloc(mem_pool);
        if(page_phyaddr == NULL) { // 若其中一块分配失败，则回滚之前分配的空间
            // 待实现...
            return NULL;
        }
        page_table_add((void*) vaddr, page_phyaddr); // 在页表中做映射
        vaddr += PG_SIZE; // 下一个虚拟页
    }
    return vaddr_start;
}

// 从内核物理内存池中申请 pg_cnt 页内存，返回虚拟地址，否则返回 NULL
void* get_kernel_pages(uint32_t pg_cnt) {
    void* vaddr = malloc_page(PF_KERNEL, pg_cnt);
    if(vaddr != NULL) { // 若分配的地址不为空，先将页框清 0，避免脏读
        memset(vaddr, 0, pg_cnt * PG_SIZE);
    }
    return vaddr;
}

// 初始化内存池
static void mem_pool_init(uint32_t all_mem) {
    put_str("  mem_pool_init start\n");
    uint32_t page_table_size = PG_SIZE * 256;       // 页表大小
    uint32_t used_mem = page_table_size + 0x100000; // 已使用的内存，0x100000 位低端1M
    uint32_t free_mem = all_mem - used_mem;         // 剩余的可用内存
    uint16_t all_free_pages = free_mem / PG_SIZE;   // 剩余的空闲内存可以分配成多少页

    // 计算内核和用户内存池可以分配到的物理容量，这里采用对半分
    uint16_t kernel_free_pages = all_free_pages / 2;
    uint16_t user_free_pages = all_free_pages - kernel_free_pages;

    // 因为位图每位表示 4k 即一页，而 1byte = 8bit，所以这里除 8 求得位图长度
    uint32_t kbm_length = kernel_free_pages / 8;
    uint32_t ubm_length = user_free_pages / 8;
    
    // 计算内核和用户内存池的起始地址
    uint32_t kp_start = used_mem; // Kernel Pool Start 内核内存池的起始地址
    uint32_t up_start = kp_start + kernel_free_pages * PG_SIZE; // User Pool Start 用户内存池的起始地址

    /* ----------------- 下面都是内核和用户内存池的初始化工作 --------------- */
    kernel_pool.phy_addr_start = kp_start;
    user_pool.phy_addr_start = up_start;

    kernel_pool.pool_size = kernel_free_pages * PG_SIZE;
    user_pool.pool_size = user_free_pages * PG_SIZE;

    kernel_pool.pool_bitmap.btmp_bytes_len = kbm_length;
    user_pool.pool_bitmap.btmp_bytes_len = ubm_length;

    /*位图是全局数据，长度不固定。
      全局或静态的数组需要在编译时直到其长度
      而我们需要根据总内存大小算出需要多少字节
      所以改为指定一块内存来生成位图*/
    kernel_pool.pool_bitmap.bits = (void*) MEM_BITMAP_BASE;
    // 用户内存池的位图紧跟在内核内存池的位图之后
    user_pool.pool_bitmap.bits = (void*) (MEM_BITMAP_BASE + kbm_length);

    /* ------------------------- 输出内存池信息 ---------------------------- */
    put_str("      kernel_pool_bitmap_start:");put_int((int)kernel_pool.pool_bitmap.bits);
    put_str(" kernel_pool_phy_addr_start:");put_int(kernel_pool.phy_addr_start);
    put_str("\n");
    put_str("      user_pool_bitmap_start:");put_int((int)user_pool.pool_bitmap.bits);
    put_str(" user_pool_phy_addr_start:");put_int(user_pool.phy_addr_start);
    put_str("\n");

    // 将位图置为 0
    bitmap_init(&kernel_pool.pool_bitmap);
    bitmap_init(&user_pool.pool_bitmap);

    /* ---------- 下面初始化内核虚拟地址的位图，按实际物理内存大小生成数组 ---------- */
    // 用于维护内核堆的虚拟地址，所以要和内核内存池大小一致
    kernel_vaddr.vaddr_bitmap.btmp_bytes_len = kbm_length;

    // 位图的数组指向一块未使用的内存，目前定位在内核内存池和用户内存池之外
    kernel_vaddr.vaddr_bitmap.bits = (void*) (MEM_BITMAP_BASE + kbm_length + ubm_length);

    kernel_vaddr.vaddr_start = K_HEAP_START;
    bitmap_init(&kernel_vaddr.vaddr_bitmap);
    put_str("   mem_pool_init done\n");
}

// 内存管理部分的初始化入口
void mem_init() {
    put_str("mem_init start\n");

    // 0xb00 是 loader.S 中定义的 mem_bytes_total 存储总内存容量
    uint32_t mem_bytes_total = (*(uint32_t*)(0xb00)); 

    mem_pool_init(mem_bytes_total);	  // 初始化内存池
    put_str("mem_init done\n");
}
